{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "czPOReizfDMY"
      },
      "source": [
        "##### Copyright 2024 Google LLC."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "4Uf6JMEbfGV-"
      },
      "outputs": [],
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "cellView": "form",
        "id": "FUqzNst0YN9P"
      },
      "outputs": [],
      "source": [
        "# The non-source code materials on this page are licensed under Creative Commons - Attribution-ShareAlike CC-BY-SA 4.0,\n",
        "# https://creativecommons.org/licenses/by-sa/4.0/legalcode."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oJt4OFKPgO8C"
      },
      "source": [
        "# Search re-ranking using Gemini embeddings"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "k5n0sCc_fKcB"
      },
      "source": [
        "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://ai.google.dev/docs/search_reranking_using_embeddings\"><img src=\"https://ai.google.dev/static/site-assets/images/docs/notebook-site-button.png\" height=\"32\" width=\"32\" />View on Google AI</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://colab.research.google.com/github/google/generative-ai-docs/blob/main/site/en/docs/search_reranking_using_embeddings.ipynb\"><img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" />Run in Google Colab</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://github.com/google/generative-ai-docs/blob/main/site/en/docs/search_reranking_using_embeddings.ipynb\"><img src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" />View source on GitHub</a>\n",
        "  </td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0UO0Cgq0mRlG"
      },
      "source": [
        "This notebook demonstrates the use of embeddings to re-rank search results. This walkthrough will focus on the following objectives:\n",
        "\n",
        "\n",
        "\n",
        "1.   Setting up your development environment and API access to use Gemini.\n",
        "2.   Using Gemini's function calling support to access the Wikipedia API.\n",
        "3.   Embedding content via Gemini API.\n",
        "4.   Re-ranking the search results.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-e5Aje-DTQOq"
      },
      "source": [
        "This is how you will implement search re-ranking:\n",
        "\n",
        "\n",
        "1.   User will query the model.\n",
        "2.   You will use Wikipedia API to return relevant search results.\n",
        "3.   The search results will be embedded and their relevance will be evaluated by calculating distance metrics like cosine similarity, dot product, etc.\n",
        "4.   Most relevant result will be returned as the final answer."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "x32Uf6GYDSDg"
      },
      "source": [
        "## Prerequisites\n",
        "\n",
        "You can run this quickstart in [Google Colab](https://colab.research.google.com/github/google/generative-ai-docs/blob/main/site/en/docs/search_reranking_using_embeddings.ipynb), which runs this notebook directly in the browser and does not require additional environment configuration.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pwImlRg9DUp7"
      },
      "source": [
        "## Setup\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0aj1YD9lDW_Z"
      },
      "source": [
        "The Python SDK for the Gemini API, is contained in the [`google-generativeai`](https://pypi.org/project/google-generativeai/) package. You will also need to install the **Wikipedia** API.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Nv_6f_qNHe_2"
      },
      "outputs": [],
      "source": [
        "!pip install -q google-generativeai"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5999faf240a4"
      },
      "outputs": [],
      "source": [
        "!pip install -q wikipedia"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "739f0bb73f05"
      },
      "source": [
        "Note: The [`wikipedia` package](https://pypi.org/project/wikipedia/) notes that it was \"designed for ease of use and simplicity, not for advanced use\", and that production or heavy use should instead \"use [Pywikipediabot](http://www.mediawiki.org/wiki/Manual:Pywikipediabot) or one of the other more advanced [Python MediaWiki API wrappers](http://en.wikipedia.org/wiki/Wikipedia:Creating_a_bot#Python)\"."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZqPr1ViPDhti"
      },
      "source": [
        "Import the necessary packages."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 34,
      "metadata": {
        "id": "D0qOtLI3FVid"
      },
      "outputs": [],
      "source": [
        "import json\n",
        "import textwrap\n",
        "\n",
        "import google.generativeai as genai\n",
        "\n",
        "import wikipedia\n",
        "from wikipedia.exceptions import DisambiguationError, PageError\n",
        "\n",
        "import numpy as np\n",
        "\n",
        "from IPython.display import Markdown\n",
        "\n",
        "def to_markdown(text):\n",
        "  text = text.replace('•', '  *')\n",
        "  return Markdown(textwrap.indent(text, '> ', predicate=lambda _: True))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qeDtGeqxDois"
      },
      "source": [
        "### Grab an API key\n",
        "\n",
        "Before you can use the Gemini API, you must first obtain an API key. If you don't already have one, create a key with one click in Google AI Studio.\n",
        "\n",
        "<a class=\"button button-primary\" href=\"https://makersuite.google.com/app/apikey\" target=\"_blank\" rel=\"noopener noreferrer\">Get an API key</a>\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PeFewbQoDqfb"
      },
      "source": [
        "In Colab, add the key to the secrets manager under the \"🔑\" in the left panel. Give it the name `GOOGLE_API_KEY`."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aeo_9J9iDtCx"
      },
      "source": [
        "Once you have the API key, pass it to the SDK. You can do this in two ways:\n",
        "\n",
        "* Put the key in the `GOOGLE_API_KEY` environment variable (the SDK will automatically pick it up from there).\n",
        "* Pass the key to `genai.configure(api_key=...)`\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 17,
      "metadata": {
        "id": "yxzEH0GqLshI"
      },
      "outputs": [],
      "source": [
        "try:\n",
        "    from google.colab import userdata\n",
        "    GOOGLE_API_KEY=userdata.get('GOOGLE_API_KEY')\n",
        "except ImportError:\n",
        "    import os\n",
        "    GOOGLE_API_KEY = os.environ['GOOGLE_API_KEY']\n",
        "\n",
        "genai.configure(api_key=GOOGLE_API_KEY)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wpVM0W-5m9fc"
      },
      "source": [
        "Key Point: Next, you will choose a model. Any embedding model will work for this tutorial, but for real applications it's important to choose a specific model and stick with it. The outputs of different models are not compatible with each other.\n",
        "\n",
        "**Note**: At this time, the Gemini API is [only available in certain regions](https://developers.generativeai.google/available_regions)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0hkJSFIWI7G1"
      },
      "source": [
        "## Define tools"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7HjN6hsAJkOq"
      },
      "source": [
        "As stated earlier, this tutorial uses Gemini's function calling support to access the Wikipedia API. Please refer to the [docs](https://ai.google.dev/docs/function_calling) to learn more about function calling."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fpFkazXUenkK"
      },
      "source": [
        "### Define the search function"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2fp-m-5Ke0XY"
      },
      "source": [
        "To cater to the search engine needs, you will design this function in the following way:\n",
        "\n",
        "\n",
        "*   For each search query, the search engine will use the `wikipedia.search` method to get relevant topics.\n",
        "*   From the relevant topics, the engine will choose `n_topics(int)` top candidates and will use `gemini-pro` to extract relevant information from the page.\n",
        "*   The engine will avoid duplicate entries by maintaining a search history.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 41,
      "metadata": {
        "id": "2U27tdP8OBys"
      },
      "outputs": [],
      "source": [
        "def wikipedia_search(search_queries: list[str]) -> list[str]:\n",
        "  \"\"\"Search wikipedia for each query and summarize relevant docs.\"\"\"\n",
        "  n_topics=3\n",
        "  search_history = set() # tracking search history\n",
        "  search_urls = []\n",
        "  mining_model = genai.GenerativeModel('gemini-pro')\n",
        "  summary_results = []\n",
        "\n",
        "  for query in search_queries:\n",
        "    print(f'Searching for \"{query}\"')\n",
        "    search_terms = wikipedia.search(query)\n",
        "\n",
        "    print(f\"Related search terms: {search_terms[:n_topics]}\")\n",
        "    for search_term in search_terms[:n_topics]: # select first `n_topics` candidates\n",
        "      if search_term in search_history: # check if the topic is already covered\n",
        "        continue\n",
        "\n",
        "      print(f'Fetching page: \"{search_term}\"')\n",
        "      search_history.add(search_term) # add to search history\n",
        "\n",
        "      try:\n",
        "        # extract the relevant data by using `gemini-pro` model\n",
        "        page = wikipedia.page(search_term, auto_suggest=False)\n",
        "        url = page.url\n",
        "        print(f\"Information Source: {url}\")\n",
        "        search_urls.append(url)\n",
        "        page = page.content\n",
        "        response = mining_model.generate_content(textwrap.dedent(f\"\"\"\\\n",
        "            Extract relevant information\n",
        "            about user's query: {query}\n",
        "            From this source:\n",
        "\n",
        "            {page}\n",
        "\n",
        "            Note: Do not summarize. Only Extract and return the relevant information\n",
        "        \"\"\"))\n",
        "\n",
        "        urls = [url]\n",
        "        if response.candidates[0].citation_metadata:\n",
        "          extra_citations = response.candidates[0].citation_metadata.citation_sources\n",
        "          extra_urls = [source.url for source in extra_citations]\n",
        "          urls.extend(extra_urls)\n",
        "          search_urls.extend(extra_urls)\n",
        "          print(\"Additional citations:\", response.candidates[0].citation_metadata.citation_sources)\n",
        "        try:\n",
        "          text = response.text\n",
        "        except ValueError:\n",
        "          pass\n",
        "        else:\n",
        "          summary_results.append(text + \"\\n\\nBased on:\\n  \" + ',\\n  '.join(urls))\n",
        "\n",
        "      except DisambiguationError:\n",
        "        print(f\"\"\"Results when searching for \"{search_term}\" (originally for \"{query}\")\n",
        "        were ambiguous, hence skipping\"\"\")\n",
        "\n",
        "      except PageError:\n",
        "        print(f'{search_term} did not match with any page id, hence skipping.')\n",
        "\n",
        "  print(f\"Information Sources:\")\n",
        "  for url in search_urls:\n",
        "    print('    ', url)\n",
        "\n",
        "  return summary_results\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 42,
      "metadata": {
        "id": "EWKkkKDXmzOX"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Searching for \"What are LLMs?\"\n",
            "Related search terms: ['Large language model', 'Prompt engineering', 'Language model']\n",
            "Fetching page: \"Large language model\"\n",
            "Information Source: https://en.wikipedia.org/wiki/Large_language_model\n",
            "Fetching page: \"Prompt engineering\"\n",
            "Information Source: https://en.wikipedia.org/wiki/Prompt_engineering\n",
            "Fetching page: \"Language model\"\n",
            "Information Source: https://en.wikipedia.org/wiki/Language_model\n",
            "Information Sources:\n",
            "     https://en.wikipedia.org/wiki/Large_language_model\n",
            "     https://en.wikipedia.org/wiki/Prompt_engineering\n",
            "     https://en.wikipedia.org/wiki/Language_model\n"
          ]
        }
      ],
      "source": [
        "example = wikipedia_search([\"What are LLMs?\"])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8bPT114QEPgW"
      },
      "source": [
        "Here is what the search results look like:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 43,
      "metadata": {
        "id": "8wy7YUFqEYTv"
      },
      "outputs": [
        {
          "data": {
            "text/markdown": [
              "> **Relevant information about LLMs:**\n",
              "> \n",
              "> * LLMs are language models notable for their ability to achieve general-purpose language generation and understanding.\n",
              "> * LLMs are artificial neural networks, the largest and most capable of which are built with a decoder-only transformer-based architecture.\n",
              "> * LLMs can be used for text generation, a form of generative AI, by taking an input text and repeatedly predicting the next token or word.\n",
              "> * Some notable LLMs are OpenAI's GPT series of models, Google's PaLM and Gemini, Meta's LLaMA family of open-source models, and Anthropic's Claude models.\n",
              "> * LLMs are trained using statistical relationships from text documents during a computationally intensive self-supervised and semi-supervised training process.\n",
              "> * LLMs are thought to acquire knowledge about syntax, semantics, and \"ontology\" inherent in human language corpora, but also inaccuracies and biases present in the corpora.\n",
              "> * LLMs can be used for a variety of tasks, including text generation, language translation, question answering, and summarization.\n",
              "> * LLMs have a number of advantages over traditional language models, including their ability to handle longer sequences of text, their ability to learn from unlabeled data, and their ability to generate more coherent and fluent text.\n",
              "> * LLMs also have a number of limitations, including their tendency to hallucinate facts, their lack of common sense knowledge, and their potential for bias.\n",
              "> * LLMs are still under development, but they have the potential to revolutionize a wide range of industries, including natural language processing, customer service, and education.\n",
              "> \n",
              "> Based on:\n",
              ">   https://en.wikipedia.org/wiki/Large_language_model"
            ],
            "text/plain": [
              "<IPython.core.display.Markdown object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "text/markdown": [
              "> **LLMs** (Large Language Models) are powerful AI models that can understand and generate text. They are designed to perform various language-related tasks, such as answering questions, summarizing documents, translating languages, writing different forms of text, and generating code. LLMs have been significantly improved through techniques such as in-context learning, which allows them to temporarily learn from specific prompts.\n",
              "> \n",
              "> **Prompt engineering** involves structuring text prompts to optimize the performance of an LLM. Effective prompts can guide the model's reasoning, provide context, and specify the desired output. Various prompt engineering techniques have been developed, including chain-of-thought prompting, generated knowledge prompting, and complexity-based prompting, each tailored to specific tasks and models.\n",
              "> \n",
              "> LLMs have also been adapted to generate images and videos. **Text-to-image** models like DALL-E 2 create art based on textual descriptions, while **text-to-video** models generate videos from textual prompts. These models require specialized prompting techniques that account for their unique capabilities and limitations.\n",
              "> \n",
              "> **Non-text prompts** are also used to guide LLMs. **Image prompting** allows users to provide images or image-based information as input, while **gradient descent**-based techniques enable the optimization of soft prompt tokens to enhance model performance.\n",
              "> \n",
              "> **Prompt injection** is a security concern where malicious users craft prompts to trick LLMs into performing unintended actions or revealing sensitive information. Mitigation strategies include input filtering, output filtering, and prompt engineering techniques to separate user input from instructions.\n",
              "> \n",
              "> LLMs are continually evolving, and new techniques and applications are being developed. They have the potential to revolutionize various industries by automating language-related tasks and enabling novel forms of creativity and communication.\n",
              "> \n",
              "> Based on:\n",
              ">   https://en.wikipedia.org/wiki/Prompt_engineering"
            ],
            "text/plain": [
              "<IPython.core.display.Markdown object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "text/markdown": [
              "> - Language models are probabilistic models of natural language.\n",
              "> - Large language models are a combination of larger datasets, feedforward neural networks, and transformers.\n",
              "> - Large language models are useful for a variety of tasks, including speech recognition, machine translation, natural language generation, optical character recognition, handwriting recognition, grammar induction, and information retrieval.\n",
              "> - Evaluation of the quality of language models is mostly done by comparison to human-created sample benchmarks created from typical language-oriented tasks.\n",
              "> \n",
              "> Based on:\n",
              ">   https://en.wikipedia.org/wiki/Language_model"
            ],
            "text/plain": [
              "<IPython.core.display.Markdown object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "from IPython.display import display\n",
        "\n",
        "for e in example:\n",
        "  display(to_markdown(e))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SKcH1LICeg5Z"
      },
      "source": [
        "### Pass the tools to the model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qQcxR7cjLjGM"
      },
      "source": [
        "If you pass a list of functions to the `GenerativeModel`'s `tools` argument,\n",
        "it will extract a schema from the function's signature and type hints, and then pass schema along to the API calls. In response the model may return a `FunctionCall` object asking to call the function.\n",
        "\n",
        "Note: This approach only handles annotations of `AllowedTypes = int | float | str | dict | list['AllowedTypes']`\n",
        "\n",
        "The `GenerativeModel` will keep a reference to the function inself, so that it _can_ execute the function locally later."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 45,
      "metadata": {
        "id": "39kqTnBRLDeQ"
      },
      "outputs": [],
      "source": [
        "model = genai.GenerativeModel(\n",
        "    'gemini-pro',\n",
        "    tools=[wikipedia_search],\n",
        "    generation_config={'temperature': 0.6})"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VedPbpzlh6jX"
      },
      "source": [
        "## Generate supporting search queries"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5sLKZik7isBW"
      },
      "source": [
        "In order to have multiple supporting search queries to the user's original query, you will ask the model to generate more such queries. This would help the engine to cover the asked question on comprehensive levels."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 46,
      "metadata": {
        "id": "Z-Ym3H9KIosY"
      },
      "outputs": [],
      "source": [
        "instructions = \"\"\"You have access to the Wikipedia API which you will be using\n",
        "to answer a user's query. Your job is to generate a list of search queries which\n",
        "might answer a user's question. Be creative by using various key-phrases from\n",
        "the user's query. To generate variety of queries, ask questions which are\n",
        "related to  the user's query that might help to find the answer. The more\n",
        "queries you generate the better are the odds of you finding the correct answer.\n",
        "Here is an example:\n",
        "\n",
        "user: Tell me about Cricket World cup 2023 winners.\n",
        "\n",
        "function_call: wikipedia_search(['What is the name of the team that\n",
        "won the Cricket World Cup 2023?', 'Who was the captain of the Cricket World Cup\n",
        "2023 winning team?', 'Which country hosted the Cricket World Cup 2023?', 'What\n",
        "was the venue of the Cricket World Cup 2023 final match?', 'Cricket World cup 2023',\n",
        "'Who lifted the Cricket World Cup 2023 trophy?'])\n",
        "\n",
        "The search function will return a list of article summaries, use these to\n",
        "answer the  user's question.\n",
        "\n",
        "Here is the user's query: {query}\n",
        "\"\"\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3Wyn3lV-d5S_"
      },
      "source": [
        "In order to yield creative and a more random variety of questions, you will set the model's temperature parameter to a value higher. Values can range from [0.0,1.0], inclusive. A value closer to 1.0 will produce responses that are more varied and creative, while a value closer to 0.0 will typically result in more straightforward responses from the model.\n",
        "\n",
        "> Note: Explore more parameters from [genai.GenerativeModel.GenerationConfig](https://ai.google.dev/api/python/google/generativeai/GenerationConfig) class to control model's response in a better way."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TD1vU6FXhYO5"
      },
      "source": [
        "## Enable automatic function calling and call the API"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qyUSnXJ5hg6x"
      },
      "source": [
        "Now start a new chat with `enable_automatic_function_calling=True`. With it enabled, the `genai.ChatSession` will handle the back and forth required to call the function, and return the final response:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 47,
      "metadata": {
        "id": "ZNJqA60yKNpT"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Searching for \"How does deep-sea life survive?\"\n",
            "Related search terms: ['Deep sea', 'Deep-sea community', 'Hydrothermal vent']\n",
            "Fetching page: \"Deep sea\"\n",
            "Information Source: https://en.wikipedia.org/wiki/Deep_sea\n",
            "Fetching page: \"Deep-sea community\"\n",
            "Information Source: https://en.wikipedia.org/wiki/Deep-sea_community\n",
            "Fetching page: \"Hydrothermal vent\"\n",
            "Information Source: https://en.wikipedia.org/wiki/Hydrothermal_vent\n",
            "Searching for \"What adaptations have deep-sea life developed to survive?\"\n",
            "Related search terms: ['Deep sea', 'Deep-sea fish', 'Deep-sea community']\n",
            "Fetching page: \"Deep-sea fish\"\n",
            "Information Source: https://en.wikipedia.org/wiki/Deep-sea_fish\n",
            "Searching for \"What are the unique characteristics of deep-sea life?\"\n",
            "Related search terms: ['Deep-sea fish', 'Deep-sea community', 'Deep sea']\n",
            "Searching for \"What are the challenges deep-sea life faces?\"\n",
            "Related search terms: ['Deep sea', 'Marine habitat', 'Sea']\n",
            "Fetching page: \"Marine habitat\"\n",
            "Information Source: https://en.wikipedia.org/wiki/Marine_habitat\n",
            "Fetching page: \"Sea\"\n",
            "Information Source: https://en.wikipedia.org/wiki/Sea\n",
            "Searching for \"How has deep-sea life evolved to cope with the extreme conditions?\"\n",
            "Related search terms: ['Deep-sea fish', 'Marine life', 'Hydrothermal vent microbial communities']\n",
            "Fetching page: \"Marine life\"\n",
            "Information Source: https://en.wikipedia.org/wiki/Marine_life\n",
            "Fetching page: \"Hydrothermal vent microbial communities\"\n",
            "Information Source: https://en.wikipedia.org/wiki/Hydrothermal_vent_microbial_communities\n",
            "Information Sources:\n",
            "     https://en.wikipedia.org/wiki/Deep_sea\n",
            "     https://en.wikipedia.org/wiki/Deep-sea_community\n",
            "     https://en.wikipedia.org/wiki/Hydrothermal_vent\n",
            "     https://en.wikipedia.org/wiki/Deep-sea_fish\n",
            "     https://en.wikipedia.org/wiki/Marine_habitat\n",
            "     https://en.wikipedia.org/wiki/Sea\n",
            "     https://en.wikipedia.org/wiki/Marine_life\n",
            "     https://en.wikipedia.org/wiki/Hydrothermal_vent_microbial_communities\n"
          ]
        }
      ],
      "source": [
        "model = genai.GenerativeModel(\n",
        "    'gemini-pro', tools=[wikipedia_search], generation_config={'temperature': 0.6})\n",
        "\n",
        "chat = model.start_chat(enable_automatic_function_calling=True)\n",
        "\n",
        "query = \"Explain how deep-sea life survives.\"\n",
        "\n",
        "res = chat.send_message(instructions.format(query=query))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 48,
      "metadata": {
        "id": "1l8KWb13M_lJ"
      },
      "outputs": [
        {
          "data": {
            "text/markdown": [
              "> Deep-sea life has evolved remarkable adaptations to survive the extreme conditions of the deep ocean. They have adapted to withstand high pressure, cold temperatures, and low oxygen levels. They have also developed unique ways to find food and communicate in the darkness of the deep sea.\n",
              "> \n",
              "> Some of the adaptations of deep-sea life include:\n",
              "> \n",
              "> * **High pressure tolerance:** Deep-sea organisms have evolved strong bodies to withstand the immense pressure of the deep ocean. Their bodies are often filled with a gelatinous substance that helps them to withstand the pressure.\n",
              "> * **Cold tolerance:** Deep-sea organisms have adapted to the cold temperatures of the deep ocean. They have enzymes that function at low temperatures, and their bodies are often covered in a thick layer of insulation.\n",
              "> * **Low oxygen tolerance:** Deep-sea organisms have adapted to the low oxygen levels of the deep ocean. They have evolved efficient respiratory systems that allow them to extract oxygen from the water.\n",
              "> * ** Bioluminescence:** Many deep-sea organisms produce their own light, a process called bioluminescence. They use bioluminescence to attract prey, communicate with each other, and defend themselves from predators.\n",
              "> * **Chemosynthesis:** Some deep-sea organisms do not rely on sunlight for food. Instead, they use a process called chemosynthesis to create food from chemicals in the water.\n",
              "> \n",
              "> These are just a few of the adaptations that deep-sea life has evolved to survive in the extreme conditions of the deep ocean. These adaptations are a testament to the resilience and adaptability of life on Earth."
            ],
            "text/plain": [
              "<IPython.core.display.Markdown object>"
            ]
          },
          "execution_count": 48,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "to_markdown(res.text)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ua7DPlj7HZyO"
      },
      "source": [
        "Check for additional citations:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 49,
      "metadata": {
        "id": "ASotq7EcHAeo"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "'No citations found'"
            ]
          },
          "execution_count": 49,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "res.candidates[0].citation_metadata or 'No citations found'"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "A0CfKMjYFSQm"
      },
      "source": [
        "That looks like it worked. You can go through the chat history to see the details of what was sent and received in the function calls:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 52,
      "metadata": {
        "id": "_9h2b8saNnxh"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "user -> {\n",
            "  \"text\": \"You have access to the Wikipedia API which you will be using\\nto answer a user's query. Your job is to generate a list of search queries which\\nmight answer a user's question. Be creative by using various key-phrases from\\nthe user's query. To generate variety of queries, ask questions which are\\nrelated to  the user's query that might help to find the answer. The more\\nqueries you generate the better are the odds of you finding the correct answer.\\nHere is an example:\\n\\nuser: Tell me about Cricket World cup 2023 winners.\\n\\nfunction_call: wikipedia_search(['What is the name of the team that\\nwon the Cricket World Cup 2023?', 'Who was the captain of the Cricket World Cup\\n2023 winning team?', 'Which country hosted the Cricket World Cup 2023?', 'What\\nwas the venue of the Cricket World Cup 2023 final match?', 'Cricket World cup 2023',\\n'Who lifted the Cricket World Cup 2023 trophy?'])\\n\\nThe search function will return a list of article summaries, use these to\\nanswer the  user's question.\\n\\nHere is the user's query: Explain how deep-sea life survives.\\n\"\n",
            "}\n",
            "------------------------------------------------------------\n",
            "model -> {\n",
            "  \"function_call\": {\n",
            "    \"name\": \"wikipedia_search\",\n",
            "    \"args\": {\n",
            "      \"search_queries\": [\n",
            "        \"How does deep-sea life survive?\",\n",
            "        \"What adaptations have deep-sea life developed to survive?\",\n",
            "        \"What are the unique characteristics of deep-sea life?\",\n",
            "        \"What are the challenges deep-sea life faces?\",\n",
            "        \"How has deep-sea life evolved to cope with the extreme conditions?\"\n",
            "      ]\n",
            "    }\n",
            "  }\n",
            "}\n",
            "------------------------------------------------------------\n",
            "user -> {\n",
            "  \"function_response\": {\n",
            "    \"name\": \"wikipedia_search\",\n",
            "    \"response\": {\n",
            "      \"result\": [\n",
            "        \"**Environmental Characteristics**\\n- Pressure: Pressure increases by about 1 atmosphere for every 10 meters of depth. Deep-sea organisms must have adaptations to withstand this pressure.\\n- Salinity: Salinity is remarkably constant throughout the deep sea, with no significant ecological differences.\\n- Temperature: The two areas of greatest temperature gradient are the transition zone between the surface waters and the deep waters (the thermocline) and the transition between the deep-sea floor and the hot water flows at the hydrothermal vents.\\n- Light: Natural light does not penetrate the deep ocean, except for the upper parts of the mesopelagic. Organisms must rely on energy sources from elsewhere, such as organic material drifting down from the photic zone.\\n\\n**Biology**\\n- Regions below the epipelagic are divided into further zones: bathyal zone (200-3000 meters), abyssal zone (3000-6000 meters), and hadal zone (6000-11,000 meters).\\n- Food: Deep-sea organisms rely on falling organic matter known as 'marine snow' and carcasses derived from the productive zone above.\\n- Adaptations: Deep-sea organisms have various adaptations to survive in extreme conditions, including: \\n 1. Jelly-like flesh to provide buoyancy\\n 2. Floaters filled with ammonium chloride that are lighter than the surrounding water\\n 3. Small size, slow metabolism, and elongated bodies\\n 4. Enhanced eyesight, such as larger eyes and rod cells for detecting light in low-light conditions\\n 5. Bioluminescence for camouflage and attracting prey\\n 6. Modifications in proteins, anatomical structures and metabolic systems to cope with high hydrostatic pressure\\n\\n**Chemosynthesis**\\n- Some species in the deep sea do not rely on dissolved organic matter for their food.\\n- These species form communities around hydrothermal vents and rely on chemosynthesis, a process where bacteria use chemical energy to produce organic matter. The tube worm Riftia is an example of an organism that benefits from this process.\\n\\n**Adaptation to Hydrostatic Pressure:**\\n- Deep-sea organisms have developed unique adaptations to survive hydrostatic pressure.\\n- Proteins can be affected by hydrostatic pressure, so deep-sea organisms have specific substitutions in the active sites of proteins like actin.\\n- These substitutions allow for better stabilization in ATP binding and subunit arrangement.\\n- Osmolytes like Trimethylamine N-oxide (TMAO) are adjusted in deep-sea fish to assist in protein stabilization.\\n- Molecular adaptations include modified Osteocalcin genes, which lead to open skulls and cartilage-based bone formation in species like the Mariana hadal snailfish. These adaptations are crucial for withstanding high pressure in the deep sea.\\n\\nBased on:\\n  https://en.wikipedia.org/wiki/Deep_sea\",\n",
            "        \"Deep-sea life survives due to various adaptations and energy sources:\\n\\n**Adaptations:**\\n- Size: Smaller size to withstand pressure.\\n- Gelatinous flesh and minimal skeletal structure.\\n- Elimination of excess cavities to prevent collapse.\\n- Eyes adapted for low light conditions.\\n- Tolerance to cold temperatures and low oxygen levels.\\n\\n**Energy Sources:**\\n\\n**Marine Snow:**\\n- Repackaged organic matter that sinks quickly, providing food for bottom-dwelling organisms.\\n\\n**Whale Falls:**\\n- Dead whales provide a significant amount of organic matter, supporting a diverse community of scavengers and other organisms.\\n- Stages of whale fall progression: mobile scavenger, opportunistic, and sulfophilic.\\n\\n**Chemosynthesis:**\\n\\n**Hydrothermal Vents:**\\n- Spew forth chemicals that bacteria can transform into energy.\\n- Support giant tube worms and other unique species.\\n- Entire ecosystems independent from sunlight.\\n\\n**Cold Seeps:**\\n- Hydrogen sulfide, methane, and other hydrocarbon-rich fluids provide energy for chemosynthetic organisms.\\n\\nBased on:\\n  https://en.wikipedia.org/wiki/Deep-sea_community\",\n",
            "        \"**How does deep-sea life survive?**\\n\\n* Life around hydrothermal vents is based on chemosynthesis, where organisms use chemical compounds as energy sources instead of sunlight.\\n* Chemosynthetic bacteria form the base of the food chain and support diverse organisms.\\n* Specialized adaptations allow organisms to withstand extreme conditions, such as high temperatures and pressure, and toxic chemicals.\\n* They have symbiotic relationships with chemoautotrophic microbial symbionts that convert inorganic molecules into organic molecules for nutrition.\\n* Their metabolism allows them to survive in environments where sunlight is absent and oxygen is limited.\\n\\nBased on:\\n  https://en.wikipedia.org/wiki/Hydrothermal_vent\",\n",
            "        \"**Adaptations of Deep-Sea Fish:**\\n\\n* **Vision**:\\n    * Large, sensitive eyes for low-light environments\\n    * Bioluminescence to attract prey or illuminate the area\\n\\n* **Sensory adaptations**:\\n    * Enhanced sensitivity to pressure and smell\\n    * Loss of eyesight in some species\\n\\n* **Buoyancy control**:\\n    * Reduction in swim bladders (in bathypelagic fish)\\n    * Hydrofoils to provide lift\\n    * High fat content and low bone density to reduce buoyancy\\n\\n* **Metabolic adaptations**:\\n    * Slow metabolism\\n    * Increased proportion of unsaturated fatty acids in cell membranes for fluidity\\n\\n* **Adaptations to high pressure**:\\n    * Gelatinous layer for buoyancy\\n    * Modifications in protein structure and reaction criteria\\n    * Rigid proteins to resist pressure\\n    * High tolerance of Na+/K+-ATPase to hydrostatic pressure\\n\\n* **Feeding adaptations**:\\n    * Long feelers to locate prey\\n    * Large mouths with sharp teeth for consuming large prey\\n    * Expandable bodies to accommodate large prey items\\n\\n* **Mating and reproduction**:\\n    * Bioluminescence to attract mates\\n    * Hermaphroditism in some species\\n    * Extreme sexual dimorphism in anglerfish (male attached to female)\\n\\nBased on:\\n  https://en.wikipedia.org/wiki/Deep-sea_fish\",\n",
            "        \"**Challenges deep-sea life faces:**\\n\\n- Extreme water pressure\\n- No sunlight\\n- Cold temperatures\\n- Limited food resources\\n- Pollution\\n- Human activities (e.g., fishing, mining)\\n\\nBased on:\\n  https://en.wikipedia.org/wiki/Marine_habitat\",\n",
            "        \"**Challenges deep-sea life faces:**\\n\\n* **Low light:** sunlight only penetrates the top 200 meters, making it difficult for plants to grow and limiting the food available for other organisms.\\n* **High pressure:** the pressure increases with depth, making it difficult for organisms to maintain their body structure and function.\\n* **Cold temperatures:** the temperature decreases with depth, making it difficult for organisms to regulate their body temperature.\\n* **Low oxygen levels:** the oxygen content of the water decreases with depth, making it difficult for organisms to breathe.\\n* **Nutrient scarcity:** the availability of nutrients decreases with depth, making it difficult for organisms to find food.\\n* **Pollution:** pollutants from human activities can accumulate in the deep sea, harming organisms and disrupting the ecosystem.\\n* **Climate change:** climate change is altering the conditions in the deep sea, such as temperature, acidity, and oxygen levels, which can be harmful to organisms.\\n\\nBased on:\\n  https://en.wikipedia.org/wiki/Sea\",\n",
            "        \"Deep-sea life has evolved to cope with the extreme conditions of the deep ocean. There is no sunlight, so primary producers must use chemosynthesis to create food. The water is cold and dark, so animals must adapt to the low temperatures and lack of light. The pressure is immense, so animals must develop strong bodies to withstand the crushing force.\\nOne of the most striking adaptations of deep-sea life is the use of bioluminescence. This is the ability to produce light, and it is used by many deep-sea animals to attract prey, communicate with each other, and defend themselves from predators. Bioluminescence is often produced by a chemical reaction that involves a luciferase enzyme and a luciferin substrate.\\nAnother adaptation of deep-sea life is the use of gigantism. This is the tendency for deep-sea animals to be larger than their shallow-water counterparts. Gigantism is thought to be an adaptation to the low food availability in the deep sea. Larger animals have a greater chance of finding food, and they can also store more energy in their bodies.\\nDeep-sea life is a fascinating and diverse group of organisms that have evolved to cope with the extreme conditions of the deep ocean. These animals have developed a range of adaptations that allow them to survive and thrive in this unique environment.\\n\\nBased on:\\n  https://en.wikipedia.org/wiki/Marine_life\",\n",
            "        \"**Adaptations** \\n\\n- Microbes that inhabit hydrothermal vents have adapted to extreme conditions, such as high temperatures, pressure, and chemical concentrations. \\n\\n- Hyperthermophiles, microorganisms that grow at temperatures above 90 \\u00b0C, are found where fluids from the vents are expelled and mixed with the surrounding water. \\n\\n- Hyperthermophilic microbes are thought to contain proteins that have extended stability at higher temperatures due to intramolecular interactions. \\n\\n- Microbes are also found in symbiotic relationships with other organisms in the hydrothermal vent environment due to their ability to have a detoxification mechanism that allows them to metabolize the sulfide-rich waters which would otherwise be toxic to the organisms and the microbes.\\n\\nBased on:\\n  https://en.wikipedia.org/wiki/Hydrothermal_vent_microbial_communities\"\n",
            "      ]\n",
            "    }\n",
            "  }\n",
            "}\n",
            "------------------------------------------------------------\n",
            "model -> {\n",
            "  \"text\": \"Deep-sea life has evolved remarkable adaptations to survive the extreme conditions of the deep ocean. They have adapted to withstand high pressure, cold temperatures, and low oxygen levels. They have also developed unique ways to find food and communicate in the darkness of the deep sea.\\n\\nSome of the adaptations of deep-sea life include:\\n\\n* **High pressure tolerance:** Deep-sea organisms have evolved strong bodies to withstand the immense pressure of the deep ocean. Their bodies are often filled with a gelatinous substance that helps them to withstand the pressure.\\n* **Cold tolerance:** Deep-sea organisms have adapted to the cold temperatures of the deep ocean. They have enzymes that function at low temperatures, and their bodies are often covered in a thick layer of insulation.\\n* **Low oxygen tolerance:** Deep-sea organisms have adapted to the low oxygen levels of the deep ocean. They have evolved efficient respiratory systems that allow them to extract oxygen from the water.\\n* ** Bioluminescence:** Many deep-sea organisms produce their own light, a process called bioluminescence. They use bioluminescence to attract prey, communicate with each other, and defend themselves from predators.\\n* **Chemosynthesis:** Some deep-sea organisms do not rely on sunlight for food. Instead, they use a process called chemosynthesis to create food from chemicals in the water.\\n\\nThese are just a few of the adaptations that deep-sea life has evolved to survive in the extreme conditions of the deep ocean. These adaptations are a testament to the resilience and adaptability of life on Earth.\"\n",
            "}\n",
            "------------------------------------------------------------\n"
          ]
        }
      ],
      "source": [
        "for content in chat.history:\n",
        "  part = content.parts[0]\n",
        "\n",
        "  print(f'{content.role} -> ', end='')\n",
        "  print(json.dumps(type(part).to_dict(part), indent=2))\n",
        "  print('---' * 20)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WDWc9Fj9Ig6A"
      },
      "source": [
        "In the chat history you can see all 4 steps:\n",
        "\n",
        "1. The user sent the query.\n",
        "2. The model replied with a `genai.protos.FunctionCall` calling the `wikipedia_search` with a number of relevant searches.\n",
        "3. Because you set `enable_automatic_function_calling=True` when creating the `genai.ChatSession`, it  executed the search function and returned the list of article summaries to the model.\n",
        "4. Folliwing the instructions in the prompt, the model generated a final answer based on those summaries.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PJP1EAgfnPUA"
      },
      "source": [
        "## [Optional] Manually execute the function call"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Sa5ke2ssKl7M"
      },
      "source": [
        "If you want to understand what happened behind the scenes, this section executes the `FunctionCall` manually to demonstrate."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 53,
      "metadata": {
        "id": "wavbHrL3K5vo"
      },
      "outputs": [],
      "source": [
        "chat = model.start_chat()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 54,
      "metadata": {
        "id": "ON4LTcmiLs2E"
      },
      "outputs": [],
      "source": [
        "result = chat.send_message(instructions.format(query=query))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "__JN3YuHe7fR"
      },
      "source": [
        "Initially the model returns a FunctionCall:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 65,
      "metadata": {
        "id": "Lgngdvcdi06F"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "{\n",
            "  \"name\": \"wikipedia_search\",\n",
            "  \"args\": {\n",
            "    \"search_queries\": [\n",
            "      \"How do deep-sea animals survive?\",\n",
            "      \"What are the adaptations of deep-sea creatures?\",\n",
            "      \"How do deep-sea animals cope with extreme pressure?\",\n",
            "      \"What are the unique characteristics of deep-sea organisms?\",\n",
            "      \"How do deep-sea animals find food?\",\n",
            "      \"How do deep-sea animals reproduce?\",\n",
            "      \"What are the challenges faced by deep-sea animals?\",\n",
            "      \"What is the role of deep-sea animals in the marine ecosystem?\"\n",
            "    ]\n",
            "  }\n",
            "}\n"
          ]
        }
      ],
      "source": [
        "fc = result.candidates[0].content.parts[0].function_call\n",
        "fc = type(fc).to_dict(fc)\n",
        "print(json.dumps(fc, indent=2))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 66,
      "metadata": {
        "id": "Bfp1Kqv7PE8N"
      },
      "outputs": [
        {
          "data": {
            "application/vnd.google.colaboratory.intrinsic+json": {
              "type": "string"
            },
            "text/plain": [
              "'wikipedia_search'"
            ]
          },
          "execution_count": 66,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "fc['name']"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xpD3axtDmfZb"
      },
      "source": [
        "Call the function with generated arguments to get the results."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 67,
      "metadata": {
        "id": "Ek4g-CTAPSou"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Searching for \"How do deep-sea animals survive?\"\n",
            "Related search terms: ['Deep sea', 'Marine life', 'Deep-sea fish']\n",
            "Fetching page: \"Deep sea\"\n",
            "Information Source: https://en.wikipedia.org/wiki/Deep_sea\n",
            "Fetching page: \"Marine life\"\n",
            "Information Source: https://en.wikipedia.org/wiki/Marine_life\n",
            "Fetching page: \"Deep-sea fish\"\n",
            "Information Source: https://en.wikipedia.org/wiki/Deep-sea_fish\n",
            "Searching for \"What are the adaptations of deep-sea creatures?\"\n",
            "Related search terms: ['Deep sea', 'Deep-sea community', 'Deep-sea fish']\n",
            "Fetching page: \"Deep-sea community\"\n",
            "Information Source: https://en.wikipedia.org/wiki/Deep-sea_community\n",
            "Searching for \"How do deep-sea animals cope with extreme pressure?\"\n",
            "Related search terms: ['Deep sea', 'Deep-sea fish', 'Deep-sea community']\n",
            "Searching for \"What are the unique characteristics of deep-sea organisms?\"\n",
            "Related search terms: ['Deep-sea fish', 'Deep sea', 'Deep-sea community']\n",
            "Searching for \"How do deep-sea animals find food?\"\n",
            "Related search terms: ['Deep-sea community', 'Deep-sea fish', 'Marine life']\n",
            "Searching for \"How do deep-sea animals reproduce?\"\n",
            "Related search terms: ['Marine life', 'Sea cucumber', 'Deep-water coral']\n",
            "Fetching page: \"Sea cucumber\"\n",
            "Information Source: https://en.wikipedia.org/wiki/Sea_cucumber\n",
            "Fetching page: \"Deep-water coral\"\n",
            "Information Source: https://en.wikipedia.org/wiki/Deep-water_coral\n",
            "Searching for \"What are the challenges faced by deep-sea animals?\"\n",
            "Related search terms: ['Deep sea', 'Marine life', 'Marine habitat']\n",
            "Fetching page: \"Marine habitat\"\n",
            "Information Source: https://en.wikipedia.org/wiki/Marine_habitat\n",
            "Searching for \"What is the role of deep-sea animals in the marine ecosystem?\"\n",
            "Related search terms: ['Marine ecosystem', 'Deep-sea community', 'Marine life']\n",
            "Fetching page: \"Marine ecosystem\"\n",
            "Information Source: https://en.wikipedia.org/wiki/Marine_ecosystem\n",
            "Information Sources:\n",
            "     https://en.wikipedia.org/wiki/Deep_sea\n",
            "     https://en.wikipedia.org/wiki/Marine_life\n",
            "     https://en.wikipedia.org/wiki/Deep-sea_fish\n",
            "     https://en.wikipedia.org/wiki/Deep-sea_community\n",
            "     https://en.wikipedia.org/wiki/Sea_cucumber\n",
            "     https://en.wikipedia.org/wiki/Deep-water_coral\n",
            "     https://en.wikipedia.org/wiki/Marine_habitat\n",
            "     https://en.wikipedia.org/wiki/Marine_ecosystem\n"
          ]
        }
      ],
      "source": [
        "summaries = wikipedia_search(**fc['args'])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kv4WGnG_gT3F"
      },
      "source": [
        "Now send the `FunctionResult` to the model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 68,
      "metadata": {
        "id": "bcLuieHqj9PW"
      },
      "outputs": [
        {
          "data": {
            "text/markdown": [
              "> **Deep-sea life survives by adapting to the extreme conditions of the deep ocean, including high pressure, low temperatures, and lack of light.**\n",
              "> \n",
              "> **Adaptations for survival include:**\n",
              "> \n",
              "> * **High internal pressure:** Deep-sea animals have high internal pressure that matches the external pressure, preventing them from being crushed.\n",
              "> * **Buoyancy adaptations:** Many deep-sea fish have a gelatinous layer below the skin or around the spine for buoyancy and swimming efficiency. They also have low tissue density, achieved through high fat content, reduced skeletal weight, and water accumulation, allowing them to float without a swim bladder.\n",
              "> * **Light and vision:** Deep-sea fish lack sunlight, so they rely on other senses, such as sensitivity to pressure changes and smell, for locating prey and mates. Many deep-sea fish are bioluminescent, using light to communicate, attract prey, or camouflage themselves. Some have sensitive eyes with high numbers of Rh1 genes, helping them see in low light conditions.\n",
              "> * **Feeding mechanisms:** Deep-sea fish often have large mouths and sharp teeth for consuming prey of similar or larger sizes. They use feelers to locate prey in the darkness.\n",
              "> * **Behavior:** Mesopelagic fish make vertical migrations following zooplankton prey, returning to deeper depths during the day. Bathypelagic fish are sedentary, waiting for prey to come close enough or being lured by bioluminescence. Some deep-sea fish are hermaphrodites, increasing their chances of reproduction in the sparse environment.\n",
              "> * **Physiological adaptations:** Deep-sea animals have slow metabolisms and unspecialized diets, allowing them to survive with limited food availability. Their proteins are structurally modified to withstand high pressure, ensuring enzymatic reactions and cellular processes function properly. Na+/K+ -ATPase, involved in osmoregulation, is more tolerant of pressure in deep-sea fish compared to shallow-water species."
            ],
            "text/plain": [
              "<IPython.core.display.Markdown object>"
            ]
          },
          "execution_count": 68,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "response = chat.send_message(\n",
        "    genai.protos.Content(\n",
        "      parts=[genai.protos.Part(\n",
        "          function_response = genai.protos.FunctionResponse(\n",
        "            name='wikipedia_search',\n",
        "            response={'result': summaries}\n",
        "          )\n",
        "      )]\n",
        "    )\n",
        ")\n",
        "\n",
        "to_markdown(response.text)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "L2Vfv8xpmuV1"
      },
      "source": [
        "## Re-ranking the search results"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cQ3JUJLeGGzA"
      },
      "source": [
        "Helper function to embed the content:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 69,
      "metadata": {
        "id": "mSkE7EynJBwF"
      },
      "outputs": [],
      "source": [
        "def get_embeddings(content: list[str]) -> np.ndarray:\n",
        "  embeddings = genai.embed_content('models/embedding-001', content, 'SEMANTIC_SIMILARITY')\n",
        "  embds = embeddings.get('embedding', None)\n",
        "  embds = np.array(embds).reshape(len(embds), -1)\n",
        "  return embds"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tip8ArqJf_ep"
      },
      "source": [
        "Please refer to the [embeddings guide](https://ai.google.dev/docs/embeddings_guide) for more information on embeddings."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nSPyycFuFj-_"
      },
      "source": [
        "Your next step is to define functions that you can use to calculate similarity scores between two embedding vectors. These scores will help you decide which embedding vector is the most relevant vector to the user's query.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ltbB0vDsKQtI"
      },
      "source": [
        "You will now implement **cosine similarity** as your metric. Here returned embedding vectors will be of unit length and hence their L1 norm (`np.linalg.norm()`) will be ~1. Hence, calculating **cosine similarity** is essentially same as calculating their **dot product score**."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 70,
      "metadata": {
        "id": "9iDFdzq_JWJW"
      },
      "outputs": [],
      "source": [
        "def dot_product(a: np.ndarray, b: np.ndarray):\n",
        "  return (a @ b.T)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MrF_1c_M_Hw3"
      },
      "source": [
        "### Similarity with user's query\n",
        "\n",
        "Now it's time to find the most relevant search result returned by the Wikipedia API.\n",
        "\n",
        "Use Gemini API to get embeddings for user's query and search results."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 71,
      "metadata": {
        "id": "gK9ryjftGDNe"
      },
      "outputs": [],
      "source": [
        "search_res = get_embeddings(summaries)\n",
        "embedded_query = get_embeddings([query])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2wwWq30uGRG3"
      },
      "source": [
        "Calculate similarity score:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 72,
      "metadata": {
        "id": "jWlFNYIsGV0X"
      },
      "outputs": [],
      "source": [
        "sim_value = dot_product(search_res, embedded_query)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bJW1pQQXG2w2"
      },
      "source": [
        "using `np.argmax` best candidate is selected.\n",
        "\n",
        "**Users's Input:** Explain how deep-sea life survives.\n",
        "\n",
        "**Answer:**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 73,
      "metadata": {
        "id": "8vDMDnCsG8Wn"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "In this document, there is no information about how deep-sea animals survive.\n",
            "\n",
            "Based on:\n",
            "  https://en.wikipedia.org/wiki/Marine_life\n"
          ]
        }
      ],
      "source": [
        "print(summaries[np.argmax(sim_value)])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ozn6mIFvoyJU"
      },
      "source": [
        "### Similarity with Hypothetical Document Embeddings (HyDE)\n",
        "\n",
        "Drawing inspiration from [[Gao et al](https://arxiv.org/abs/2212.10496)] the objective here is to generate a template answer to the user's query using `gemini-pro`'s internal knowledge. This hypothetical answer will serve as a baseline to calculate relevance of all the search results."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 74,
      "metadata": {
        "id": "J7m5KAkMREBH"
      },
      "outputs": [
        {
          "data": {
            "text/markdown": [
              "> In the enigmatic depths where sunlight surrenders to inky blackness, life perseveres, illuminated by the faintest of luminescent flickers.\n",
              "> \n",
              "> Imagine a realm where extreme pressure could crush the mightiest of vessels. Yet, in this unforgiving abyss, creatures have evolved with bodies resilient as the very seafloor. Their flexible exoskeletons or gelatinous tissues withstand the crushing weight gracefully.\n",
              "> \n",
              "> Oxygen, a lifeline for most creatures, grows scarce with depth. Enter our deep-sea dwellers, whose bodies have ingeniously adapted. They absorb oxygen directly through their skin or gills, maximizing every molecule they find.\n",
              "> \n",
              "> Nutrient scarcity plagues these depths, where sunlight cannot penetrate to foster photosynthesis. Instead, these creatures rely on chemosynthesis, a remarkable process that utilizes chemicals from hydrothermal vents or decaying matter.\n",
              "> \n",
              "> In the perpetual darkness, vision becomes obsolete. Instead, sensory organs have evolved to detect minute vibrations, bioluminescence, and heat gradients, guiding them through the shadowy labyrinth.\n",
              "> \n",
              "> Temperature fluctuations can be drastic, from freezing cold to scalding heat. But deep-sea creatures have mastered the art of thermoregulation, their internal systems finely tuned to withstand the extremes.\n",
              "> \n",
              "> Growth is a slow and arduous process in these unforgiving depths. Many species exhibit extreme longevity, surviving for centuries or even millennia. Their life cycles are meticulously paced, ensuring their survival in this harsh environment.\n",
              "> \n",
              "> Reproduction is a perilous task, with offspring often vulnerable and exposed. Some deep-sea creatures protect their young with parental care, nurturing them until they can fend for themselves in this unforgiving realm.\n",
              "> \n",
              "> The deep sea, a testament to the resilience and adaptability of life, is a fascinating and mysterious world. Its inhabitants continue to inspire awe and wonder, reminding us of the extraordinary diversity and ingenuity that exists within our planet's watery depths."
            ],
            "text/plain": [
              "<IPython.core.display.Markdown object>"
            ]
          },
          "execution_count": 74,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "hypothetical_ans_model = genai.GenerativeModel('gemini-pro')\n",
        "res = hypothetical_ans_model.generate_content(f\"\"\"Generate a hypothetical answer\n",
        "to the user's query by using your own knowledge. Assume that you know everything\n",
        "about the said topic. Do not use factual information, instead use placeholders\n",
        "to complete your answer. Your answer should feel like it has been written by a human.\n",
        "\n",
        "query: {query}\"\"\")\n",
        "\n",
        "to_markdown(res.text)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "85f4MpLgrYbW"
      },
      "source": [
        "Use Gemini API to get embeddings for the baseline answer and compare them with search results"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 75,
      "metadata": {
        "id": "fSI6-5eYI3me"
      },
      "outputs": [],
      "source": [
        "hypothetical_ans = get_embeddings([res.text])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WxYGEWbqrh44"
      },
      "source": [
        "Calculate similarity scores to rank the search results"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 76,
      "metadata": {
        "id": "99sjvMtlJfL_"
      },
      "outputs": [],
      "source": [
        "sim_value = dot_product(search_res, hypothetical_ans)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 77,
      "metadata": {
        "id": "C32IhIapQFNa"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "array([[0.72687077],\n",
              "       [0.73694087],\n",
              "       [0.77235092],\n",
              "       [0.75185433],\n",
              "       [0.63363508],\n",
              "       [0.62639701],\n",
              "       [0.71418557],\n",
              "       [0.70211815]])"
            ]
          },
          "execution_count": 77,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "sim_value"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wJ7pklxursSk"
      },
      "source": [
        "using `np.argmax` best candidate is selected.\n",
        "\n",
        "**Users's Input:** Explain how deep-sea life survives.\n",
        "\n",
        "**Answer:**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 78,
      "metadata": {
        "id": "MzTfyU_mJ8-M"
      },
      "outputs": [
        {
          "data": {
            "text/markdown": [
              "> **How do deep-sea animals survive?**\n",
              "> \n",
              "> **Adaptations to Pressure:**\n",
              "> \n",
              "> * Deep-sea animals have high internal pressure that matches the external pressure, preventing them from being crushed.\n",
              "> * Their cell membranes contain a higher proportion of unsaturated fatty acids, which increases membrane fluidity in high-pressure environments.\n",
              "> \n",
              "> **Buoyancy Adaptations:**\n",
              "> \n",
              "> * Many deep-sea fish have a gelatinous layer below the skin or around the spine for buoyancy and swimming efficiency.\n",
              "> * They have low tissue density, achieved through high fat content, reduced skeletal weight, and water accumulation, allowing them to float without a swim bladder.\n",
              "> \n",
              "> **Light and Vision:**\n",
              "> \n",
              "> * Deep-sea fish lack sunlight, so they rely on other senses, such as sensitivity to pressure changes and smell, for locating prey and mates.\n",
              "> * Many deep-sea fish are bioluminescent, using light to communicate, attract prey, or camouflage themselves.\n",
              "> * Some have sensitive eyes with high numbers of Rh1 genes, helping them see in low light conditions.\n",
              "> \n",
              "> **Feeding Mechanisms:**\n",
              "> \n",
              "> * Deep-sea fish often have large mouths and sharp teeth for consuming prey of similar or larger sizes.\n",
              "> * They use feelers to locate prey in the darkness.\n",
              "> \n",
              "> **Behavior:**\n",
              "> \n",
              "> * Mesopelagic fish make vertical migrations following zooplankton prey, returning to deeper depths during the day.\n",
              "> * Bathypelagic fish are sedentary, waiting for prey to come close enough or being lured by bioluminescence.\n",
              "> * Some deep-sea fish are hermaphrodites, increasing their chances of reproduction in the sparse environment.\n",
              "> \n",
              "> **Physiological Adaptations:**\n",
              "> \n",
              "> * Deep-sea animals have slow metabolisms and unspecialized diets, allowing them to survive with limited food availability.\n",
              "> * Their proteins are structurally modified to withstand high pressure, ensuring enzymatic reactions and cellular processes function properly.\n",
              "> * Na+/K+ -ATPase, involved in osmoregulation, is more tolerant of pressure in deep-sea fish compared to shallow-water species.\n",
              "> \n",
              "> Based on:\n",
              ">   https://en.wikipedia.org/wiki/Deep-sea_fish"
            ],
            "text/plain": [
              "<IPython.core.display.Markdown object>"
            ]
          },
          "execution_count": 78,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "to_markdown(summaries[np.argmax(sim_value)])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UjDN2OUpsL6M"
      },
      "source": [
        "You have now created a search re-ranking engine using embeddings!"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AKH2vs2Lc1R1"
      },
      "source": [
        "## Next steps\n",
        "\n",
        "To learn how to use other services in the Gemini API, visit the [Python quickstart](https://ai.google.dev/tutorials/python_quickstart). To learn more about how you can use the embeddings, check out the [examples](https://ai.google.dev/examples?keywords=embed) available."
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "search_reranking_using_embeddings.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
